import time
import os
import math
import scipy.optimize
import matplotlib.pyplot as plt
from scipy.interpolate import interp1d
import numpy

import util as ref
import refprop2 as basicRef
import convection as conv



def makePlots(data,location='.'):
    for i in range(len(data)):
         try:
            plt.clf()
            plt.grid(True)
            plt.ylabel(data[13][i])
            plt.xlabel(data[13][4])
            plt.plot(data[4],data[i],'o')  
            plt.savefig(location+os.sep+'plot'+str(i+1)+'.png')
            plt.clf()
         except:
             print ' '
    return(True)
    
def writeResults(filename,data):
    fileOut = open(filename,'w')
    for each in data[13]:
        fileOut.write(each+',')    
    fileOut.write('\n')
    for j in range(len(data[4])):
        for i in range(len(data)):
            try:
                fileOut.write(str(data[i][j])+',')
            except:
                fileOut.write(',')
        fileOut.write('\n')
    fileOut.close()  
    return(True)    

def conservBalance(x,rho1,h1,u1,p1,A1,dh,A2,dx,q,pLoss,fluid,nChan=24):
    '''input x is array of [rho2,h2,u2]'''
    #print pLoss
    massImb = (rho1*u1*A1 - x[0]*x[2]*A2)/(rho1*u1*A1)
    p2 = ref.pressure_rho_h(x[0],x[1],fluid)*1000 
    momImb = -A2*(x[0]*pow(x[2],2) + p2) + A1*(p1+rho1*pow(u1,2)) + (p1+p2)/2*(A2-A1) - pLoss*(A2+A1)/2
    momImb = momImb / (A1*(p1+rho1*pow(u1,2))+(p1+p2)/2*(A2-A1))
    energyImb = ((x[1]*1000 + 0.5*pow(x[2],2) - (h1*1000 + 0.5*pow(u1,2) + q))/(h1*1000 + 0.5*pow(u1,2) + q))
    #print massImb,momImb,energyImb
    totalImb = abs(massImb)+abs(momImb)+abs(energyImb)
    return(totalImb)

def march(p0,h0,mdot,A,X,D,hc,tHot,nC,
    res=20,t0=0.00,fluid='WATER.FLD',
    saturated=False,kwall=30.,k = 0.00001):
    if saturated:
        rho0 = ref.denThrottle_Hup_Pdown(h0,p0,fluid)
        t0 = ref.PDFLSH(p0,rho0/ref.wm.value)[0]        
    elif t0>0.000:
        props = ref.TPFLSH(t0,p0)
        h0 = props[5]/ref.wm.value
        rho0 = props[0]*ref.wm.value
    else:
        sys.exit()
    u0 = mdot/rho0/A[0]
    mu,kcoolant = basicRef.TRNPRP(t0,rho0/ref.wm.value)
    out =[ [rho0],[h0],[u0],[rho0*u0*A[0]],[X[0]],[A[0]],[D[0]],[hc[0]],[0.00],
        [p0*1000],[t0],[mu*1e-6],[kcoolant],[],[0.0],[0.0],[0.0],[0.0],[0.0],[0.0]
        ,[0.0],[0.0],[0.0]]
    out[13] = ['den(kg/m3)','enthalpy(kJ/kg)','Vel(m/s)','mdot(kg/s)','X(m)'
                ,'Achan(m2)','Dcham(m)','hg(W/m2/K)','q''(W/m2)','Pregen(Pa)'
                ,'Tregen(K)','muregen(microPa/s)','kregen(W/m/K)','names','cp'
                ,'Pr','Nu','hcool(W/m2/K)','Reynolds Number','qCalc W/m^2'
                ,'Twcold (K)','Twhot (K)','Hydraulic Diameter (m)']
    n = res
    for j in range(1,len(A)):
        nChan = nC[j]
        As = ref.myfrange(A[j-1],A[j],n)
        Xs = ref.myfrange(X[j-1],X[j],n)
        Ds = ref.myfrange(D[j-1],D[j],n)
        hs = ref.myfrange(hc[j-1],hc[j],n)
        tHots = ref.myfrange(tHot[j-1],tHot[j],n)
        for i in range(1,n,1):
            qprime = (tHots[i]-out[21][-1])*hs[i]
            dx = Xs[i]-Xs[i-1]
            dh = math.pow((As[i-1]+As[i])/2/nChan*4/math.pi,0.5)
            q = qprime*(dx)*Ds[i]*math.pi/mdot 
            Re = reynoldsNumber(out[2][-1],dh,out[11][-1]/out[0][-1])
            f = frictionFactor(Re,k,dh)
            pLoss = frictionLoss(f,dx,dh,out[0][-1],out[2][-1])
            # blah is [rho2,h2,u2]
            blah = scipy.optimize.fmin(conservBalance,
                                       [out[0][-1],out[1][-1],out[2][-1]],
                                       args=(out[0][-1],out[1][-1],out[2][-1],out[9][-1],As[i-1],dh,As[i]
                                       ,dx,q,pLoss,fluid),disp=False)#,xtol=1e-19)
            out[0].append(blah[0])
            out[1].append(blah[1])
            out[2].append(blah[2])
            out[3].append(blah[0]*blah[2]*As[i])
            out[4].append(Xs[i])
            out[5].append(As[i])
            out[6].append(Ds[i])
            out[7].append(hs[i])
            out[8].append(qprime)
            props = ref.DHFLSH(blah[0]/ref.wm.value,blah[1]*ref.wm.value)
            out[9].append(props[1]*1000)
            out[10].append(props[0])
            mu,kcoolant = basicRef.TRNPRP(out[10][-1],blah[0]/ref.wm.value)
            out[11].append(mu*1.0e-6)
            out[12].append(kcoolant)
            out[14].append(max(props[8],30)/ref.wm.value)
            out[15].append(out[14][-1]*mu/kcoolant/1000)
            out[16].append(0.023*math.pow(Re,0.8)*math.pow(max(0.001,out[15][-1]),0.4))
            out[17].append(kcoolant*out[16][-1]/dh)
            out[18].append(Re)
            convans = conv.convWall(out[7][-1],tHots[i],out[17][-1],out[10][-1],kwall)
            out[19].append(convans[0])
            out[20].append(convans[1])
            out[21].append(convans[2])
            out[22].append(dh)
            print '.',
        print '.'
    return(out)
    
def march_new(p0,h0,mdot,A,X,D,hhg,tHot,nC,
    res=20,t0=0.00,fluid='WATER.FLD',
    saturated=False,kwall=30.,k = 0.00001,reverse=False):
    
    Xs = numpy.linspace(X[0],X[1],res)
    
    fAofX = interp1d(A[0],A[1])
    fDofX = interp1d(D[0],D[1])
    fHcofX = interp1d(hhg[0],hhg[1])
    fThotofX = interp1d(tHot[0],tHot[1])
    fnCofX = interp1d(nC[0],nC[1])
    
    As          = fAofX(Xs)
    Ds          = fDofX(Xs)
    Hhgs         = fHcofX(Xs)
    tHots       = fThotofX(Xs)
    nCs         = fnCofX(Xs)

    if reverse:
        Xs = Xs[::-1]
        As = As[::-1]
        Ds = Ds[::-1]
        Hhgs = Hhgs[::-1]
        tHots = tHots[::-1]
        nCs = nCs[::-1]        
 
    
    if saturated:
        rho0 = ref.denThrottle_Hup_Pdown(h0,p0,fluid)
        t0 = ref.PDFLSH(p0,rho0/ref.wm.value)[0]        
    elif t0>0.000:
        props = ref.TPFLSH(t0,p0)
        h0 = props[5]/ref.wm.value
        rho0 = props[0]*ref.wm.value
    else:
        sys.exit()
    u0 = mdot/rho0/As[0]
    mu,kcoolant = basicRef.TRNPRP(t0,rho0/ref.wm.value)
    out =[ [rho0],[h0],[u0],[rho0*u0*As[0]],[Xs[0]],[As[0]],[Ds[0]],[Hhgs[0]],[0.00],
        [p0*1000],[t0],[mu*1e-6],[kcoolant],[],[0.0],[0.0],[0.0],[0.0],[0.0],[0.0]
        ,[0.0],[0.0],[0.0],[nCs[0]],[0.0]]
    out[13] = ['den(kg/m3)','enthalpy(kJ/kg)','Vel(m/s)','mdot(kg/s)','X(m)'
                ,'Achan(m2)','Dcham(m)','hg(W/m2/K)','q''(W/m2)','Pregen(Pa)'
                ,'Tregen(K)','muregen(microPa/s)','kregen(W/m/K)','names','cp'
                ,'Pr','Nu','hcool(W/m2/K)','Reynolds Number','qCalc W/m^2'
                ,'Twcold (K)','Twhot (K)','Hydraulic Diameter (m)','nChannels'
                ,'Mass Error']
    for i in range(1,len(Xs)):
   # for each in Xs:
            qprime = (tHots[i]-out[21][-1])*Hhgs[i]
            dx = abs(Xs[i]-Xs[i-1])
            dh = math.pow((As[i-1]+As[i])/2/nCs[i]*4/math.pi,0.5)
            q = qprime*(dx)*Ds[i]*math.pi/mdot 
            Re = reynoldsNumber(out[2][-1],dh,out[11][-1]/out[0][-1])
            f = frictionFactor(Re,k,dh)
            pLoss = frictionLoss(f,dx,dh,out[0][-1],out[2][-1])
            # blah is [rho2,h2,u2]
            blah = scipy.optimize.fmin(conservBalance,
                                       [out[0][-1],out[1][-1],out[2][-1]],
                                       args=(out[0][-1],out[1][-1],out[2][-1],out[9][-1],As[i-1],dh,As[i]
                                       ,dx,q,pLoss,fluid),disp=False)#,xtol=1e-19)
            out[0].append(blah[0])
            out[1].append(blah[1])
            out[2].append(blah[2])
            out[3].append(blah[0]*blah[2]*As[i])
            out[4].append(Xs[i])
            out[5].append(As[i])
            out[6].append(Ds[i])
            out[7].append(Hhgs[i])
            out[8].append(qprime)
            props = ref.DHFLSH(blah[0]/ref.wm.value,blah[1]*ref.wm.value)
            out[9].append(props[1]*1000)
            out[10].append(props[0])
            mu,kcoolant = basicRef.TRNPRP(out[10][-1],blah[0]/ref.wm.value)
            out[11].append(mu*1.0e-6)
            out[12].append(kcoolant)
            out[14].append(max(props[8],30)/ref.wm.value)
            out[15].append(out[14][-1]*mu/kcoolant/1000)
            out[16].append(0.023*math.pow(Re,0.8)*math.pow(max(0.001,out[15][-1]),0.4))
            out[17].append(kcoolant*out[16][-1]/dh)
            out[18].append(Re)
            convans = conv.convWall(out[7][-1],tHots[i],out[17][-1],out[10][-1],kwall)
            out[19].append(convans[0])
            out[20].append(convans[1])
            out[21].append(convans[2])
            out[22].append(dh)
            out[23].append(nCs[i])
            out[24].append(1-blah[0]*blah[2]*As[i]/out[3][0])
            print '.',
    print '\n'
    return(out)    

def colebrooke(f,Re,k,dh):
        leftHand = 1 / math.pow(f,0.5)
        #print f,Re,k,dh
        rightHand = -2.0*math.log10 ((2.51 / (Re*math.pow(f,0.5)) + (k / dh)/3.72))
        return (abs(leftHand - rightHand))
                                
def frictionFactor(Re,k,dh):
    '''Returns friction factor given Reynolds number, roughness, and hydraulic diameter'''
    f = scipy.optimize.fmin(colebrooke,[0.01],args=[Re,k,dh],disp=False,xtol=1e-6)
    return (f[0])
    
def frictionLoss(f,l,dh,rho,v):
    ploss = f *l / dh *math.pow(v,2) / 2.0 * rho
    return(ploss)    

def reynoldsNumber(V,dh,vis):
    '''Velocity (m/s), hydraulic diameter (m), viscosity (m^2/s)'''
    #print V,dh,vis
    Re = V*dh/vis
    return(Re)

def scaleH(h,p1,p2):
    '''scale h as (p1/p2)^0.8'''
    out=[]
    p1 = float(p1)
    p2 = float(p2)
    for each in h:
        out.append(each*math.pow(p2/p1,0.8))
    return (out)

def scale(x,v1,v2):
    out=[]
    p1 = float(v1)
    p2 = float(v2)
    for each in x:
        out.append(each*(v2/v1))
    return (out)